{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4ba51a29-a566-4b5a-97f0-10e634567e40",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "source": [
    "# MCPify your AWS Lambda\n",
    "## Transform AWS Lambda functions into secure MCP tools with Bedrock AgentCore Gateway\n",
    "\n",
    "## Overview\n",
    "Bedrock AgentCore Gateway provides customers a way to turn their existing AWS Lambda functions into fully-managed MCP servers without needing to manage infra or hosting. Gateway will provide a uniform Model Context Protocol (MCP) interface across all these tools. Gateway employs a dual authentication model to ensure secure access control for both incoming requests and outbound connections to target resources. The framework consists of two key components: Inbound Auth, which validates and authorizes users attempting to access gateway targets, and Outbound Auth, which enables the gateway to securely connect to backend resources on behalf of authenticated users. Gateways uses IAM role to authorize the calls to AWS Lambda functions for outb ound authorization.\n",
    "\n",
    "![How does it work](images/lambda-iam-gateway.png)\n",
    "\n",
    "### Tutorial Details\n",
    "\n",
    "\n",
    "| Information          | Details                                                   |\n",
    "|:---------------------|:----------------------------------------------------------|\n",
    "| Tutorial type        | Interactive                                               |\n",
    "| AgentCore components | AgentCore Gateway, AgentCore Identity                     |\n",
    "| Agentic Framework    | Strands Agents                                            |\n",
    "| Gateway Target type  | AWS Lambda                                                |\n",
    "| Inbound Auth IdP     | Amazon Cognito                                            |\n",
    "| Outbound Auth        | AWS IAM                                                   |\n",
    "| LLM model            | Anthropic Claude Sonnet 3.7, Amazon Nova Pro              |\n",
    "| Tutorial components  | Creating AgentCore Gateway and Invoking AgentCore Gateway |\n",
    "| Tutorial vertical    | Cross-vertical                                            |\n",
    "| Example complexity   | Easy                                                      |\n",
    "| SDK used             | boto3                                                     |\n",
    "\n",
    "In the first part of the tutorial we will create some AmazonCore Gateway targets\n",
    "\n",
    "### Tutorial Architecture\n",
    "In this tutorial we will transform operations defined in AWS lambda function into MCP tools and host it in Bedrock AgentCore Gateway.\n",
    "For demonstration purposes, we will use a Strands Agent using Amazon Bedrock models\n",
    "In our example we will use a very simple agent with two tools: get_order and update_order."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5362e1ad-f027-4452-a8d9-0b861c0115c2",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "To execute this tutorial you will need:\n",
    "* Jupyter notebook (Python kernel)\n",
    "* uv\n",
    "* AWS credentials\n",
    "* Amazon Cognito"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e93da982-0d11-427c-9958-fff3c80b32f1",
   "metadata": {},
   "source": [
    "## Configuring Authentication for Incoming AgentCore Gateway Requests\n",
    "AgentCore Gateway provides secure connections via inbound and outbound authentication. For the inbound authentication, the AgentCore Gateway analyzes the OAuth token passed during invocation to decide allow or deny the access to a tool in the gateway. If a tool needs access to external resources, the AgentCore Gateway can use outbound authentication via API Key, IAM or OAuth Token to allow or deny the access to the external resource.\n",
    "\n",
    "\n",
    "\n",
    "During the inbound authorization flow, an agent or the MCP client calls an MCP tool in the AgentCore Gateway adding an OAuth access token (generated from the user’s IdP). AgentCore Gateway then validates the OAuth access token and performs inbound authorization.\n",
    "\n",
    "If the tool running in AgentCore Gateway needs to access external resources, OAuth will retrieve credentials of downstream resources using the resource credential provider for the Gateway target. AgentCore Gateway pass the authorization credentials to the caller to get access to the downstream API. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e326677b-11b5-4488-b01b-573ad275a764",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure you download the latest botocore and boto3 libraries.\n",
    "import shutil\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "def ensure_uv_installed():\n",
    "    if shutil.which(\"uv\") is None:\n",
    "        print(\"🔧 'uv' not found. Installing with pip...\")\n",
    "        subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"uv\"])\n",
    "    else:\n",
    "        print(\"✅ 'uv' is already installed.\")\n",
    "\n",
    "def uv_install(*packages):\n",
    "    ensure_uv_installed()\n",
    "    uv_path = shutil.which(\"uv\")\n",
    "    print(f\"📦 Installing {', '.join(packages)} using uv...\")\n",
    "    subprocess.check_call([uv_path, \"pip\", \"install\", *packages])\n",
    "\n",
    "uv_install(\"botocore\", \"boto3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4e6248d-b740-418b-ae0d-c0a9623e43e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set AWS credentials if not using Amazon SageMaker notebook\n",
    "import os\n",
    "os.environ['AWS_ACCESS_KEY_ID'] = '' # set the access key\n",
    "os.environ['AWS_SECRET_ACCESS_KEY'] = '' # set the secret \n",
    "os.environ['AWS_DEFAULT_REGION'] = 'us-east-1' # set the AWS region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd5f3851-d0ca-484f-bb41-68f66ddbea44",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "# Get the directory of the current script\n",
    "if '__file__' in globals():\n",
    "    current_dir = os.path.dirname(os.path.abspath(__file__))\n",
    "else:\n",
    "    current_dir = os.getcwd()  # Fallback if __file__ is not defined (e.g., Jupyter)\n",
    "\n",
    "# Navigate to the directory containing utils.py (one level up)\n",
    "utils_dir = os.path.abspath(os.path.join(current_dir, '..'))\n",
    "\n",
    "# Add to sys.path\n",
    "sys.path.insert(0, utils_dir)\n",
    "\n",
    "# Now you can import utils\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddf8bfe9-8435-44cc-9e8a-ed8ece48f6a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Create a sample AWS Lambda function that you want to convert into MCP tools\n",
    "lambda_resp = utils.create_gateway_lambda(\"lambda_function_code.zip\")\n",
    "\n",
    "if lambda_resp is not None:\n",
    "    if lambda_resp['exit_code'] == 0:\n",
    "        print(\"Lambda function created with ARN: \", lambda_resp['lambda_function_arn'])\n",
    "    else:\n",
    "        print(\"Lambda function creation failed with message: \", lambda_resp['lambda_function_arn'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef17018-d2be-40e3-a892-55dd420297f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Create an IAM role for the Gateway to assume\n",
    "import utils\n",
    "agentcore_gateway_iam_role = utils.create_agentcore_gateway_role(\"sample-lambdagateway\")\n",
    "print(\"Agentcore gateway role ARN: \", agentcore_gateway_iam_role['Role']['Arn'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65120594-c3ec-4d51-810b-8d478851d8d2",
   "metadata": {},
   "source": [
    "# Create Amazon Cognito Pool for Inbound authorization to Gateway"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bbbc736-6f5e-4a60-a598-ead43a0c1f90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating Cognito User Pool \n",
    "import os\n",
    "import boto3\n",
    "import requests\n",
    "import time\n",
    "from botocore.exceptions import ClientError\n",
    "\n",
    "REGION = os.environ['AWS_DEFAULT_REGION']\n",
    "USER_POOL_NAME = \"sample-agentcore-gateway-pool\"\n",
    "RESOURCE_SERVER_ID = \"sample-agentcore-gateway-id\"\n",
    "RESOURCE_SERVER_NAME = \"sample-agentcore-gateway-name\"\n",
    "CLIENT_NAME = \"sample-agentcore-gateway-client\"\n",
    "SCOPES = [\n",
    "    {\"ScopeName\": \"gateway:read\", \"ScopeDescription\": \"Read access\"},\n",
    "    {\"ScopeName\": \"gateway:write\", \"ScopeDescription\": \"Write access\"}\n",
    "]\n",
    "scopeString = f\"{RESOURCE_SERVER_ID}/gateway:read {RESOURCE_SERVER_ID}/gateway:write\"\n",
    "\n",
    "cognito = boto3.client(\"cognito-idp\", region_name=REGION)\n",
    "\n",
    "print(\"Creating or retrieving Cognito resources...\")\n",
    "user_pool_id = utils.get_or_create_user_pool(cognito, USER_POOL_NAME)\n",
    "print(f\"User Pool ID: {user_pool_id}\")\n",
    "\n",
    "utils.get_or_create_resource_server(cognito, user_pool_id, RESOURCE_SERVER_ID, RESOURCE_SERVER_NAME, SCOPES)\n",
    "print(\"Resource server ensured.\")\n",
    "\n",
    "client_id, client_secret  = utils.get_or_create_m2m_client(cognito, user_pool_id, CLIENT_NAME, RESOURCE_SERVER_ID)\n",
    "print(f\"Client ID: {client_id}\")\n",
    "\n",
    "# Get discovery URL  \n",
    "cognito_discovery_url = f'https://cognito-idp.{REGION}.amazonaws.com/{user_pool_id}/.well-known/openid-configuration'\n",
    "print(cognito_discovery_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1a63450-7fb9-42fc-ab4f-3d86c27bb2f8",
   "metadata": {},
   "source": [
    "# Create the Gateway with Amazon Cognito Authorizer for inbound authorization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72f2cd57-7777-42d3-b6f3-c45ed0a935c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# CreateGateway with Cognito authorizer without CMK. Use the Cognito user pool created in the previous step\n",
    "gateway_client = boto3.client('bedrock-agentcore-control', region_name = os.environ['AWS_DEFAULT_REGION'])\n",
    "auth_config = {\n",
    "    \"customJWTAuthorizer\": { \n",
    "        \"allowedClients\": [client_id],  # Client MUST match with the ClientId configured in Cognito. Example: 7rfbikfsm51j2fpaggacgng84g\n",
    "        \"discoveryUrl\": cognito_discovery_url\n",
    "    }\n",
    "}\n",
    "create_response = gateway_client.create_gateway(name='TestGWforLambda',\n",
    "    roleArn = agentcore_gateway_iam_role['Role']['Arn'], # The IAM Role must have permissions to create/list/get/delete Gateway \n",
    "    protocolType='MCP',\n",
    "    authorizerType='CUSTOM_JWT',\n",
    "    authorizerConfiguration=auth_config, \n",
    "    description='AgentCore Gateway with AWS Lambda target type'\n",
    ")\n",
    "print(create_response)\n",
    "# Retrieve the GatewayID used for GatewayTarget creation\n",
    "gatewayID = create_response[\"gatewayId\"]\n",
    "gatewayURL = create_response[\"gatewayUrl\"]\n",
    "print(gatewayID)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94dc57be-e50e-4997-a440-3bcf582d09bc",
   "metadata": {},
   "source": [
    "# Create an AWS Lambda target and transform into MCP tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c1474da-cdd6-4dbd-84e7-3fdad018647d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace the AWS Lambda function ARN below\n",
    "lambda_target_config = {\n",
    "    \"mcp\": {\n",
    "        \"lambda\": {\n",
    "            \"lambdaArn\": lambda_resp['lambda_function_arn'], # Replace this with your AWS Lambda function ARN\n",
    "            \"toolSchema\": {\n",
    "                \"inlinePayload\": [\n",
    "                    {\n",
    "                        \"name\": \"get_order_tool\",\n",
    "                        \"description\": \"tool to get the order\",\n",
    "                        \"inputSchema\": {\n",
    "                            \"type\": \"object\",\n",
    "                            \"properties\": {\n",
    "                                \"orderId\": {\n",
    "                                    \"type\": \"string\"\n",
    "                                }\n",
    "                            },\n",
    "                            \"required\": [\"orderId\"]\n",
    "                        }\n",
    "                    },                    \n",
    "                    {\n",
    "                        \"name\": \"update_order_tool\",\n",
    "                        \"description\": \"tool to update the orderId\",\n",
    "                        \"inputSchema\": {\n",
    "                            \"type\": \"object\",\n",
    "                            \"properties\": {\n",
    "                                \"orderId\": {\n",
    "                                    \"type\": \"string\"\n",
    "                                }\n",
    "                            },\n",
    "                            \"required\": [\"orderId\"]\n",
    "                        }\n",
    "                    }\n",
    "                ]\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "credential_config = [ \n",
    "    {\n",
    "        \"credentialProviderType\" : \"GATEWAY_IAM_ROLE\"\n",
    "    }\n",
    "]\n",
    "targetname='LambdaUsingSDK'\n",
    "response = gateway_client.create_gateway_target(\n",
    "    gatewayIdentifier=gatewayID,\n",
    "    name=targetname,\n",
    "    description='Lambda Target using SDK',\n",
    "    targetConfiguration=lambda_target_config,\n",
    "    credentialProviderConfigurations=credential_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3ac6532-5299-4024-917d-bcd60caea6ed",
   "metadata": {},
   "source": [
    "# Calling Bedrock AgentCore Gateway from a Strands Agent\n",
    "\n",
    "The Strands agent seamlessly integrates with AWS tools through the Bedrock AgentCore Gateway, which implements the Model Context Protocol (MCP) specification. This integration enables secure, standardized communication between AI agents and AWS services.\n",
    "\n",
    "At its core, the Bedrock AgentCore Gateway serves as a protocol-compliant Gateway that exposes fundamental MCP APIs: ListTools and InvokeTools. These APIs allow any MCP-compliant client or SDK to discover and interact with available tools in a secure, standardized way. When the Strands agent needs to access AWS services, it communicates with the Gateway using these MCP-standardized endpoints.\n",
    "\n",
    "The Gateway's implementation adheres strictly to the (MCP Authorization specification)[https://modelcontextprotocol.org/specification/draft/basic/authorization], ensuring robust security and access control. This means that every tool invocation by the Strands agent goes through authorization step, maintaining security while enabling powerful functionality.\n",
    "\n",
    "For example, when the Strands agent needs to access MCP tools, it first calls ListTools to discover available tools, then uses InvokeTools to execute specific actions. The Gateway handles all the necessary security validations, protocol translations, and service interactions, making the entire process seamless and secure.\n",
    "\n",
    "This architectural approach means that any client or SDK that implements the MCP specification can interact with AWS services through the Gateway, making it a versatile and future-proof solution for AI agent integrations.\n",
    "\n",
    "![Strands agent calling Gateway](images/strands-lambda-gateway.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60756f03-352b-46bc-9cb6-c845bae7acc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "uv_install(\"mcp[cli]\", \"strands-agents\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "873031fe-62b5-4196-91be-500c1f87dfd4",
   "metadata": {},
   "source": [
    "# Request the access token from Amazon Cognito for inbound authorization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ed1d1d6-e84c-4286-bf25-3ad6a49723b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Requesting the access token from Amazon Cognito authorizer...May fail for some time till the domain name propogation completes\")\n",
    "token_response = utils.get_token(user_pool_id, client_id, client_secret,scopeString,REGION)\n",
    "token = token_response[\"access_token\"]\n",
    "print(\"Token response:\", token)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6fd0379d-9576-43cb-aa9b-72b86c43b472",
   "metadata": {},
   "source": [
    "# Strands agent calling MCP tools of AWS Lambda using Bedrock AgentCore Gateway"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be4b39c4-7387-4ce8-b728-f2347fbdaa36",
   "metadata": {},
   "outputs": [],
   "source": [
    "from strands.models import BedrockModel\n",
    "from mcp.client.streamable_http import streamablehttp_client \n",
    "from strands.tools.mcp.mcp_client import MCPClient\n",
    "from strands import Agent\n",
    "\n",
    "def create_streamable_http_transport():\n",
    "    return streamablehttp_client(gatewayURL,headers={\"Authorization\": f\"Bearer {token}\"})\n",
    "\n",
    "client = MCPClient(create_streamable_http_transport)\n",
    "\n",
    "## The IAM credentials configured in ~/.aws/credentials should have access to Bedrock model\n",
    "yourmodel = BedrockModel(\n",
    "    model_id=\"us.amazon.nova-pro-v1:0\",\n",
    "    temperature=0.7,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e64794e7-4e5f-4fc5-824a-61c901e356c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from strands import Agent\n",
    "import logging\n",
    "\n",
    "\n",
    "# Configure the root strands logger. Change it to DEBUG if you are debugging the issue.\n",
    "logging.getLogger(\"strands\").setLevel(logging.INFO)\n",
    "\n",
    "# Add a handler to see the logs\n",
    "logging.basicConfig(\n",
    "    format=\"%(levelname)s | %(name)s | %(message)s\", \n",
    "    handlers=[logging.StreamHandler()]\n",
    ")\n",
    "\n",
    "with client:\n",
    "    # Call the listTools \n",
    "    tools = client.list_tools_sync()\n",
    "    # Create an Agent with the model and tools\n",
    "    agent = Agent(model=yourmodel,tools=tools) ## you can replace with any model you like\n",
    "    print(f\"Tools loaded in the agent are {agent.tool_names}\")\n",
    "    print(f\"Tools configuration in the agent are {agent.tool_config}\")\n",
    "    # Invoke the agent with the sample prompt. This will only invoke  MCP listTools and retrieve the list of tools the LLM has access to. The below does not actually call any tool.\n",
    "    agent(\"Hi , can you list all tools available to you\")\n",
    "    # Invoke the agent with sample prompt, invoke the tool and display the response\n",
    "    agent(\"Check the order status for order id 123 and show me the exact response from the tool\")\n",
    "    # Call the MCP tool explicitly. The MCP Tool name and arguments must match with your AWS Lambda function or the OpenAPI/Smithy API\n",
    "    result = client.call_tool_sync(\n",
    "    tool_use_id=\"get-order-id-123-call-1\", # You can replace this with unique identifier. \n",
    "    name=targetname+\"___get_order_tool\", # This is the tool name based on AWS Lambda target types. This will change based on the target name\n",
    "    arguments={\"orderId\": \"123\"}\n",
    "    )\n",
    "    # Print the MCP Tool response\n",
    "    print(f\"Tool Call result: {result['content'][0]['text']}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73169486-2fa0-42b5-8a8f-d6bdcc532159",
   "metadata": {},
   "source": [
    "**Issue: if you get below error while executing below cell, it indicates incompatibily between pydantic and pydantic-core versions.**\n",
    "\n",
    "```\n",
    "TypeError: model_schema() got an unexpected keyword argument 'generic_origin'\n",
    "```\n",
    "**How to resolve?**\n",
    "\n",
    "You will need to make sure you have pydantic==2.7.2 and pydantic-core 2.27.2 that are both compatible. Restart the kernel once done."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa8dd91d-15a6-45b2-b539-bb3eaf0e0e08",
   "metadata": {},
   "source": [
    "# Clean up\n",
    "\n",
    "Additional resources are also created like IAM role, IAM Policies, Credentials provider, AWS Lambda functions, Cognito user pools, s3 buckets that you might need to manually delete as part of the clean up. This depends on the example you run."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a22f158d-83b6-4d65-892f-9a208af3d742",
   "metadata": {},
   "source": [
    "## Delete the gateway (Optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3edc5b6-e5f0-47b9-994f-378415bad50d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import utils\n",
    "utils.delete_gateway(gateway_client,gatewayID)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
